#include "hyper/Hyper.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeRange.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Support/TypeID.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include <iostream>
#include <ostream>

#include "crypto/Crypto.h"
#include "llvm/ADT/ArrayRef.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

#define CHECK_ATTR(op) \
  do{ \
    auto device_attr = dyn_cast_or_null<StringAttr>(op->getAttr("device")); \
    auto device_str = device_attr.getValue(); \
    if (!device_str.contains("cgra")) { \
      return failure(); \
    } \
  }while(0);


using namespace mlir;
using namespace vector;

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

func::FuncOp get_or_insert_func(Location loc, ModuleOp themodule, StringRef func_name, FunctionType func_ty) {
    func::FuncOp existingFunc = themodule.lookupSymbol<func::FuncOp>(func_name);
    if(existingFunc) {
      return existingFunc;
    }
    auto algo_run = func::FuncOp::create(loc, func_name, func_ty);
    algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
    themodule.push_back(algo_run);
    return algo_run;
}


class Sha1LoweringCGRAPattern : public OpRewritePattern<crypto::Sha1core_OP> {
public:
  using OpRewritePattern<crypto::Sha1core_OP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::Sha1core_OP op,
                                PatternRewriter &rewriter) const override {
    CHECK_ATTR(op)
    auto loc = op.getLoc();

    // create extern function sha1_run if not exist.
    StringRef algo_func_name = "cgra_sha1_run";
    auto m = op->getParentOfType<ModuleOp>();
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto algoFuncType = rewriter.getFunctionType( mlir::TypeRange{dynI8MemrefType, dynI8MemrefType, rewriter.getI32Type()}, mlir::TypeRange{});
    
    func::FuncOp algo_func = get_or_insert_func(loc, m, algo_func_name, algoFuncType);


    Value input = op.getInput();
    Value output = op.getOutput();
    auto plain_len = op.getMsgLen();
    assert(plain_len);
    if (output.getType() != dynI8MemrefType) {
      output = rewriter.create<memref::CastOp>(
          loc, dynI8MemrefType, output);
    }
    auto func_call = rewriter.create<mlir::func::CallOp>(
        loc, algo_func, mlir::ValueRange{input, output, plain_len});

    rewriter.replaceOp(op, func_call->getResults());
    return mlir::success();
  }
};



class MD5LoweringCGRAPattern : public OpRewritePattern<crypto::MD5OP> {
public:
  using OpRewritePattern<crypto::MD5OP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::MD5OP op,
                                PatternRewriter &rewriter) const override {
    CHECK_ATTR(op)
    auto loc = op.getLoc();

    // create extern function sha1_run if not exist.
    StringRef algo_func_name = "cgra_md5_run";
    auto m = op->getParentOfType<ModuleOp>();
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto algoFuncType = rewriter.getFunctionType( mlir::TypeRange{dynI8MemrefType, dynI8MemrefType, rewriter.getI32Type()}, mlir::TypeRange{});
    
    func::FuncOp algo_func = get_or_insert_func(loc, m, algo_func_name, algoFuncType);

    // call func
    Value input = op.getInput();
    Value output = op.getOutput();
    Value plain_len = op.getMsgLen();
    if (output.getType() != dynI8MemrefType) {
      output = rewriter.create<memref::CastOp>(
          loc, dynI8MemrefType, output);
    }
    auto func_call = rewriter.create<mlir::func::CallOp>(
        loc, algo_func,
        mlir::ValueRange{input, output, plain_len});
    rewriter.replaceOp(op, func_call->getResults());
    return mlir::success();
  }
};

class SM3LoweringCGRAPattern : public OpRewritePattern<crypto::SM3OP> {
public:
  using OpRewritePattern<crypto::SM3OP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::SM3OP op,
                                PatternRewriter &rewriter) const override {
    CHECK_ATTR(op)
    auto loc = op.getLoc();

    // create extern function sha1_run if not exist.
    StringRef algo_func_name = "cgra_sm3_run";
    auto m = op->getParentOfType<ModuleOp>();
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto algoFuncType = rewriter.getFunctionType( mlir::TypeRange{dynI8MemrefType, dynI8MemrefType, rewriter.getI32Type()}, mlir::TypeRange{});
    
    func::FuncOp algo_func = get_or_insert_func(loc, m, algo_func_name, algoFuncType);

    // call func
    Value input = op.getInput();
    Value output = op.getOutput();
    Value plain_len = op.getMsgLen();
    if (output.getType() != dynI8MemrefType) {
      output = rewriter.create<memref::CastOp>(loc, dynI8MemrefType, output);
    }
    auto func_call = rewriter.create<mlir::func::CallOp>(
    loc, algo_func, mlir::ValueRange{input, output, plain_len});
    rewriter.replaceOp(op, func_call->getResults());
    return mlir::success();
  }
};

class AESLoweringCGRAPattern : public OpRewritePattern<crypto::AESOP> {
public:
  using OpRewritePattern<crypto::AESOP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::AESOP op,
                                PatternRewriter &rewriter) const override {
    // CHECK_ATTR(op)
    do{ 
      auto device_attr = dyn_cast_or_null<StringAttr>(op->getAttr("device")); 
      auto device_str = device_attr.getValue(); 
      if (!device_str.contains("cgra")) {
        return failure();
      } 
    }while(0);
    auto loc = op.getLoc();

    // create extern function sha1_run if not exist.
    StringRef algo_func_name = "cgra_aes_run";
    auto m = op->getParentOfType<ModuleOp>();
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());
    auto algoFuncType = rewriter.getFunctionType(mlir::TypeRange{dynI8MemrefType, dynI8MemrefType, rewriter.getI32Type(), dynI8MemrefType}, mlir::TypeRange{});
    func::FuncOp algo_func = get_or_insert_func(loc, m, algo_func_name, algoFuncType);

    // call func
    Value input = op.getInput();
    Value output = op.getOutput();
    Value plain_len = op.getMsgLen();
    Value key = op.getKey();
    if (output.getType() != dynI8MemrefType) {
      output = rewriter.create<memref::CastOp>(
          loc, dynI8MemrefType, output);
    }
    auto func_call = rewriter.create<mlir::func::CallOp>(
        loc, algo_func, mlir::ValueRange{input, output, plain_len, key});
    rewriter.replaceOp(op, func_call->getResults());
    return mlir::success();
  }
};


} // namespace



void populateConvertCryptoToCGRAPatterns(RewritePatternSet &patterns) {
  patterns.add<Sha1LoweringCGRAPattern>(patterns.getContext());
  patterns.add<MD5LoweringCGRAPattern>(patterns.getContext());
  patterns.add<SM3LoweringCGRAPattern>(patterns.getContext());
  patterns.add<AESLoweringCGRAPattern>(patterns.getContext());
}

namespace {

class LowerCryptoToCGRAPass
    : public PassWrapper<LowerCryptoToCGRAPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerCryptoToCGRAPass)
  LowerCryptoToCGRAPass() = default;
  LowerCryptoToCGRAPass(const LowerCryptoToCGRAPass &) {}
  StringRef getArgument() const final { return "lower-crypto-to-cgra"; }
  StringRef getDescription() const final {
    return "Lower Crypto Dialect to CGRA function call.";
  }

  void runOnOperation() override;
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<func::FuncDialect>();
  }
};
} // end anonymous namespace.

void LowerCryptoToCGRAPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<func::FuncDialect, memref::MemRefDialect,
                         arith::ArithDialect>();
  target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>();

  RewritePatternSet patterns(context);
  populateConvertCryptoToCGRAPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerCryptoToCGRAPass() {
  PassRegistration<LowerCryptoToCGRAPass>();
}

std::unique_ptr<Pass> createLowerCryptoToCGRAPass() {
  return std::make_unique<LowerCryptoToCGRAPass>();
}
